package com.xiyucode.yuerinterface.controller;

import com.fasterxml.jackson.core.JsonProcessingException;
import com.fasterxml.jackson.databind.JsonNode;
import com.fasterxml.jackson.databind.ObjectMapper;
import com.zhipu.oapi.ClientV4;
import com.zhipu.oapi.Constants;
import com.zhipu.oapi.service.v4.model.ChatCompletionRequest;
import com.zhipu.oapi.service.v4.model.ChatMessage;
import com.zhipu.oapi.service.v4.model.ChatMessageRole;
import com.zhipu.oapi.service.v4.model.ModelApiResponse;
import lombok.extern.slf4j.Slf4j;
import org.springframework.web.bind.annotation.*;

import java.util.ArrayList;
import java.util.List;
import java.util.concurrent.TimeUnit;

/**
 * zhipuAI Controller
 *
 * @author xiyucode
 */
@RestController
@RequestMapping("/zhipuAi")
@Slf4j
public class ZhiPuAiController {

    private static final String requestIdTemplate = "request_1_";

    private static final ObjectMapper mapper = new ObjectMapper();

    @PostMapping("/ask")
    public String askZhiPuAI(@RequestParam String apiKey, String text){

        ClientV4 client = new ClientV4.Builder(apiKey)
                .networkConfig(30, 10, 40, 40, TimeUnit.SECONDS)
                .build();


        ObjectMapper objectMapper = new ObjectMapper();
        List<ChatMessage> messages = new ArrayList<>();
        ChatMessage chatMessage = new ChatMessage(ChatMessageRole.USER.value(), text);

        messages.add(chatMessage);
        String requestId = String.format(requestIdTemplate, System.currentTimeMillis());

        ChatCompletionRequest chatCompletionRequest = ChatCompletionRequest.builder()
                .model(Constants.ModelChatGLM4)
                .stream(Boolean.FALSE)
                .invokeMethod(Constants.invokeMethod)
                .messages(messages)
                .requestId(requestId)
                .build();
        ModelApiResponse invokeModelApiResp = client.invokeModelApi(chatCompletionRequest);
        try {
            String result = mapper.writeValueAsString(invokeModelApiResp);

            // 将所需信息分离出来
            JsonNode jsonNode = objectMapper.readTree(result);
            JsonNode contentNode = jsonNode.path("data").path("choices").path(0).path("message").path("content");
            if (!contentNode.isMissingNode()) {
                String content = contentNode.asText();
                return content;
            } else {
                return "error";
            }
        } catch (JsonProcessingException e) {
            e.printStackTrace();
        }
        return "error";
    }

}
